#!/usr/bin/env python3
import argparse
import grpc
import os
import sys
from time import sleep

sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 'utils/'))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper

def main(p4info_file_path, bmv2_file_path):
    print("=== Starting Fixed MAC Controller ===")
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1', address='127.0.0.1:50051', device_id=0)

        s1.MasterArbitrationUpdate()
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)

        print("Adding routes with correct MAC addresses...")

        # 使用正确的MAC地址（从ARP表中看到的）
        # h1 -> 端口1
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.ipv4_lpm",
            match_fields={"hdr.ipv4.dstAddr": ["10.0.1.1", 32]},
            action_name="MyIngress.ipv4_forward",
            action_params={"dstAddr": "00:00:00:00:01:01", "port": 1}  # 修正MAC
        )
        s1.WriteTableEntry(table_entry)
        print("✓ Route for h1 (10.0.1.1) -> port 1")

        # h2 -> 端口2
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.ipv4_lpm",
            match_fields={"hdr.ipv4.dstAddr": ["10.0.1.2", 32]},
            action_name="MyIngress.ipv4_forward",
            action_params={"dstAddr": "00:00:00:00:01:02", "port": 2}  # 修正MAC
        )
        s1.WriteTableEntry(table_entry)
        print("✓ Route for h2 (10.0.1.2) -> port 2")

        # h3 -> 端口3
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.ipv4_lpm",
            match_fields={"hdr.ipv4.dstAddr": ["10.0.1.3", 32]},
            action_name="MyIngress.ipv4_forward",
            action_params={"dstAddr": "00:00:00:00:01:03", "port": 3}  # 修正MAC
        )
        s1.WriteTableEntry(table_entry)
        print("✓ Route for h3 (10.0.1.3) -> port 3")

        # h4 -> 端口4
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.ipv4_lpm",
            match_fields={"hdr.ipv4.dstAddr": ["10.0.1.4", 32]},
            action_name="MyIngress.ipv4_forward",
            action_params={"dstAddr": "00:00:00:00:01:04", "port": 4}  # 修正MAC
        )
        s1.WriteTableEntry(table_entry)
        print("✓ Route for h4 (10.0.1.4) -> port 4")

        print("\n🎉 All routes added with correct MAC addresses!")
        print("Now test ping in Mininet...")

        while True:
            sleep(10)

    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info file',
                        type=str, default='./build/acl.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file',
                        type=str, default='./build/acl.json')
    args = parser.parse_args()
    main(args.p4info, args.bmv2_json)
